import pymongo
import random

from utils.log import logger
from settings import MONGO_URL
from domain import Proxy


class MongoPool(object):
    """
    代理池数据库模块
    """
    def __init__(self):
        # 链接mongodb数据库
        self.client = pymongo.MongoClient(MONGO_URL)
        # 获取要操作的集合
        self.proxies = self.client['proxy_pool']['proxies']

    def __del__(self):
        # 关闭数据库的链接
        self.client.close()

    def insert(self, proxy=None):
        """保存代理ip到数据库中"""
        count = self.proxies.count_documents({'_id': proxy.ip})
        # 如果值不为0, 说明该代理IP, 已经存在了, 直接返回
        if count == 0:
            dic = proxy.__dict__
            dic['_id'] = proxy.ip
            self.proxies.insert_one(dic)
            # 记录日志
            logger.info('插入新代理ip：{}'.format(proxy))
        else:
            logger.warning('已经存在代理ip：{}'.format(proxy))
            # for i in self.find({'_id': proxy.ip}):
            #     print('已经存在代理ip：{}'.format(i))
            # self.update(proxy)

    def update(self, proxy=None):
        """更新代理"""
        if proxy:
            self.proxies.update_one({"_id": proxy.ip}, {"$set": proxy.__dict__})
            logger.info("更新代理：{}".format(proxy))
        else:
            logger.error("请传入要更新的代理")

    def insert_or_update(self, proxy):
        if proxy:
            count = self.proxies.count_documents({'_id': proxy.ip})
            # 如果值不为0, 说明该代理IP, 已经存在了, 直接返回
            if count == 0:
                dic = proxy.__dict__
                dic['_id'] = proxy.ip
                self.proxies.insert_one(dic)
                # 记录日志
                logger.info('插入新代理ip：{}'.format(proxy))
            else:
                self.proxies.update_one({"_id": proxy.ip}, {"$set": proxy.__dict__})
                logger.info("更新代理：{}".format(proxy))
        else:
            logger.error("请传入要更新的代理")

    def delete_one(self, proxy):
        """删除代理"""
        self.proxies.delete_one({"_id": proxy.ip})
        logger.info("删除代理ip：{}".format(proxy))

    def find_all(self):
        """查询所有代理ip"""
        cursor = self.proxies.find()
        for item in cursor:
            # Proxy没有“_id”，需要删除这个key
            item.pop("_id")
            proxy = Proxy(**item)
            yield proxy

    def find(self, conditions=None, count=0):
        """
        实现查询功能，根据条件查询，查询指定数量
        :param conditions: 查询条件字典
        :param count: 限制最多取出多少代理ip
        :return: 返回满足条件的代理ip（proxy对象）列表
        """
        cursor = self.proxies.find(conditions, limit=count).sort(
            [("score", pymongo.DESCENDING), ("speed", pymongo.ASCENDING)])

        proxy_list = []
        for item in cursor:
            item.pop('_id')
            proxy = Proxy(**item)
            proxy_list.append(proxy)

        return proxy_list

    def get_proxies(self, protocol=None, domain=None, count=0):
        """
        根据协议类型，和要访问网站的域名, 获取代理ip列表
        :param protocol: 协议：http, https
        :param domain: 域名： jd.com
        :param count: 用于限制获取多少个代理ip，默认所有
        :param nick_type: 匿名类型：默认高匿代理ip
        :return: 满足要求的代理ip
        """
        # 定义查询条件
        conditions = {}
        # 根据协议， 指定查询条件
        if protocol is None:
            # 若没有传入协议类型，返回两种都支持的
            conditions['protocol'] = 2
        elif protocol.lower() == 'http':  # lower()转小写
            conditions['protocol'] = {'$in': [2, 0]}
        elif protocol.lower() == 'https':
            conditions['protocol'] = {'$in': [2, 1]}

        if domain:
            conditions['disable_domains'] = {'$nin': [domain]}

        return self.find(conditions=conditions, count=count)

    def random_proxy(self, protocol=None, domain=None, count=0):
        """根据协议类型，和要访问网站的域名, 随机获取一个代理ip"""
        try:
            proxy = self.get_proxies(protocol=protocol, domain=domain, count=count)
            return random.choice(proxy)
        except Exception as e:
            logger.error("没有获到默认类型【高匿】代理")
            return "没有获到该类型代理"

    def disable_domain(self, ip, domain):
        """
        实现把指定域名添加到指定ip的disable_domain列表中
        :param ip: IP地址
        :param domain: 域名
        :return: True or False
        """
        if self.proxies.count_documents({"_id": ip, 'disable_domains': domain}) == 0:
            self.proxies.update_one({'_id': ip}, {'$push': {'disable_domains': domain}})
            return True
        return False

if __name__ == "__main__":
    mongo = MongoPool()

    dic = {'area': None, 'speed': 5.97, 'port': '38094', 'nick_type': 2, 'protocol': 1, 'ip': '222.135.92.67', 'disable_domains': ['jd.com'], 'score': 30}

    # mongo.delete_one(proxy)

    # proxy = Proxy(**dic)
    # mongo.update(proxy)
    # mongo.delete_one(proxy)
    # mongo.insert(proxy)
    # mongo.insert_or_update(proxy)

    # for proxy in mongo.find_all():
    #     print(proxy)

    for proxy in mongo.find({"ip":'113.108.242.36'}):
        print(proxy)

    # for proxy in mongo.get_proxies(protocol='http', nick_type=2):
    #     print(proxy)

    # for proxy in mongo.random_proxy(protocol='http'):
    #     print(proxy)

    # print(mongo.random_proxy(protocol='https'))

    # print(mongo.disable_domain('222.135.92.66', 'taobao.com'))

    # for proxy in mongo.find({"ip":'113.108.242.36'}):
    #         print(proxy)
